modelopt/onnx/quantization/qdq_utils.py (1)

928-930: Use op_type-based detection for Constant nodes
Replace the substring match on "Constant" in reshape_node.input with an explicit check of the producer node’s op_type to avoid relying on name patterns. For example:

- shape_constant_name = next(i for i in reshape_node.input if "Constant" in i)
+ shape_constant_name = next(
+     i for i in reshape_node.input
+     if tensor_producer_map[i].op_type == "Constant"
+ )

tests/unit/onnx/test_qdq_utils.py (4)

70-76: Redundant Cast between Reshape and Transpose (optional)

dq_output is already FLOAT; inserting weight_cast to FLOAT is a no-op unless it’s intentionally exercising cast-conversion logic. If this Cast is only for testing that path, consider a brief comment or renaming the node to indicate intent; otherwise, remove it to keep the pattern minimal (DQ → Reshape → Transpose).

Also applies to: 80-80

---

96-96: Avoid seeding an unused scale initializer when constant_scale=True

When constant_scale=True, the graph still includes a scale initializer up-front, which can make the “new scale initializer” assertion trivially pass. Suggest emitting it conditionally so the test precisely validates the pass rewires the scale.

Apply this diff:

-    nodes = [dq_node, reshape_constant, reshape_node, cast_node, transpose_node, matmul_node]
+    nodes = [dq_node, reshape_constant, reshape_node, cast_node, transpose_node, matmul_node]
@@
-        initializer=[weight_tensor, scale_tensor],
+        initializer=[weight_tensor] if constant_scale else [weight_tensor, scale_tensor],

Also applies to: 104-104

---

252-257: Name-based Cast preservation: tighten the match (nit)

if "norm/Cast" in node.name may over-match unrelated nodes. Consider anchoring (e.g., regex (^|/)layer_norm/Cast$) to avoid false positives as models grow.

---

315-316: FP4 tests look consistent with packed 4-bit along axis 0; add a negative/edge test

The 2D inputs and uint8 expectations align with _cast_fp4 packing the first dimension by 2. Add a failure case for odd first-dimension to lock the contract, and optionally a 4D case to exercise batching.

Example additions:

def test_cast_fp4_odd_first_dim_raises():
    with pytest.raises(AssertionError):
        _cast_fp4(np.zeros((3, 2), dtype=np.float32))

def test_cast_fp4_4d_batch():
    x = np.random.randn(2, 2, 2, 2).astype(np.float32)  # first dim even
    y = _cast_fp4(x)
    assert y.dtype == np.uint8
    assert y.shape[0] == 1 and y.shape[1:] == x.shape[1:]

Also applies to: 320-322, 325-327, 330-332, 335-337, 340-342, 348-348

examples/onnx_ptq/torch_quant_to_onnx.py (3)

86-92: Avoid heavy model instantiation just to read input size (optional)

Creating a full pretrained model here can be slow. If available in your timm version, prefer fetching the default cfg to get input_size without instantiating weights; otherwise, at least consider pretrained=False for this helper.

---

122-127: Add basic validation for --batch_size (and align with data size) (optional)

Help text implies constraints but they’re not enforced. Guard against invalid values and optionally round calibration_data_size to full batches for determinism.

Example (place after args = parser.parse_args()):

if args.batch_size <= 0:
    raise ValueError("--batch_size must be > 0")
if args.calibration_data_size <= 0:
    raise ValueError("--calibration_data_size must be > 0")
# Optional: ensure full batches only
# args.calibration_data_size = (args.calibration_data_size // args.batch_size) * args.batch_size

---

141-149: DataLoader returns GPU tensors from worker processes (risk of pickling/IPC issues)

load_calibration_data constructs CUDA tensors and then uses num_workers=4. Multiprocessing with GPU tensors can be brittle and memory-heavy. Prefer keeping tensors on CPU in the loader and moving them to device in the forward loop, or set num_workers=0 for simplicity.

Suggested adjustments:

• Keep calibration tensors on CPU in load_calibration_data, and enable pin_memory=True:

return torch.utils.data.DataLoader(
    calib_tensor, batch_size=batch_size, shuffle=True, num_workers=0, pin_memory=True
)

• Move batches to device inside forward_loop:

def forward_loop(model):
    for batch in data_loader:
        if isinstance(batch, torch.Tensor):
            batch = batch.to(device, non_blocking=True)
        model(batch)

• Optionally set shuffle=False for reproducible calibration.

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

modelopt/onnx/quantization/qdq_utils.py (1)

614-627: FP4 packing along wrong axis; assertion is on axis 0 instead of last axis; docstring missing requirement

FP4 should be packed along the last dimension (matches NVFP4QTensor.quantize). Asserting and reshaping on dim 0 will break many weight shapes. Also document the even-length requirement.

-def _cast_fp4(array: np.ndarray) -> np.ndarray:
-    """Cast a numpy array to FLOAT4E2M1 using PyTorch."""
-    array_f32_t = torch.from_numpy(array)
-    array_f32_t_shape = array_f32_t.shape
-    assert array_f32_t_shape[0] % 2 == 0, "array_f32_t_shape[0] must be divisible by 2"
-    array_f4_t_shape = (array_f32_t_shape[0] // 2, *array_f32_t_shape[1:])
-    if torch.cuda.is_available():
-        array_f32_t = array_f32_t.cuda()
-    array_f4_t = NVFP4QTensor._cast_fp4(array_f32_t)
-    array_f4_t = array_f4_t.flatten()
-    array_f4_t_packed = (array_f4_t[::2] | (array_f4_t[1::2] << 4)).reshape(array_f4_t_shape)
-    array_f4 = array_f4_t_packed.cpu().numpy().astype(np.uint8)
-    return array_f4
+def _cast_fp4(array: np.ndarray) -> np.ndarray:
+    """Cast a numpy array to FLOAT4E2M1 using PyTorch.
+
+    Note: The last dimension must be even; two FP4 values are packed per byte.
+    """
+    array_f32_t = torch.from_numpy(array)
+    if array_f32_t.shape[-1] % 2 != 0:
+        raise ValueError(
+            f"Last dimension must be divisible by 2 for FP4 packing; got {array_f32_t.shape[-1]}"
+        )
+    if torch.cuda.is_available():
+        array_f32_t = array_f32_t.cuda()
+    q4 = NVFP4QTensor._cast_fp4(array_f32_t)  # values in [0..15], same shape as input
+    packed = (q4[..., 1::2] << 4) | q4[..., 0::2]  # pack along last dim
+    return packed.cpu().numpy().astype(np.uint8)   # shape: (*, last_dim//2)

setup.py (1)

50-51: ONNX/ORT version alignment check

Bumping to onnx~=1.19.0 and onnxconverter-common~=1.16.0 looks fine. Please verify ORT packages still resolve and support new dtypes (Float4/Float8) used elsewhere. Also consider aligning onnxruntime-directml to ~=1.22.0 for consistency unless there’s a known constraint.

examples/onnx_ptq/torch_quant_to_onnx.py (3)

86-92: Input shape: validate batch_size > 0

Return shape logic is good. Add a guard for batch_size >= 1 to avoid silent mis-shapes.

 def get_model_input_shape(model_name, batch_size):
     """Get the input shape from timm model configuration."""
     model = timm.create_model(model_name, pretrained=True, num_classes=1000)
     data_config = timm.data.resolve_model_data_config(model)
     input_size = data_config["input_size"]
-    return (batch_size, *tuple(input_size))  # Add batch dimension
+    if batch_size < 1:
+        raise ValueError(f"batch_size must be >= 1, got {batch_size}")
+    return (batch_size, *tuple(input_size))  # Add batch dimension

---

122-127: CLI: constrain batch_size

Add argparse-level validation and clarify help.

 parser.add_argument(
     "--batch_size",
-    type=int,
-    default=1,
-    help="Batch size for calibration.",
+    type=int,
+    default=1,
+    choices=range(1, 1024),
+    metavar="{1..1023}",
+    help="Batch size for calibration (>=1).",
 )

---

131-150: Minor: set eval/inference during quantization

Set model.eval() before quantize and rely on inference_mode in forward_loop (as suggested above) to speed up and stabilize calibration behavior.

-    # Quantize model
-    quantized_model = quantize_model(model, config, data_loader)
+    model.eval()
+    quantized_model = quantize_model(model, config, data_loader)

modelopt/onnx/quantization/qdq_utils.py (4)

951-956: Assumes Cast always follows Reshape

Guard against missing/extra nodes to avoid crashes; gracefully skip if pattern doesn’t match.

-        # Remove unnecessary Cast node
-        cast_node = reshape_child_nodes[0]
-        assert cast_node.op_type == "Cast", f"Expected Cast node for {node.name}"
-        nodes_to_remove.append(cast_node.name)
-        cast_child_nodes = [n for n in graph.node if cast_node.output[0] in n.input]
+        # Remove unnecessary Cast node (if present)
+        cast_child_nodes = reshape_child_nodes
+        if reshape_child_nodes and reshape_child_nodes[0].op_type == "Cast":
+            cast_node = reshape_child_nodes[0]
+            nodes_to_remove.append(cast_node.name)
+            cast_child_nodes = [n for n in graph.node if cast_node.output[0] in n.input]

---

958-977: Be robust if there is no Transpose after Cast

Avoid indexing without checks.

-        if cast_child_nodes[0].op_type == "Transpose":
-            transpose_node = cast_child_nodes[0]
+        if cast_child_nodes and cast_child_nodes[0].op_type == "Transpose":
+            transpose_node = cast_child_nodes[0]
             nodes_to_remove.append(transpose_node.name)
             ...
-        else:
-            matmul_node = cast_child_nodes[0]
+        else:
+            assert cast_child_nodes, f"No consumer found after Cast/Reshape for {node.name}"
+            matmul_node = cast_child_nodes[0]

---

1007-1021: Pre-quant scale cleanup assumes exactly one Cast child

Loosen the assumptions: skip if multiple/no children or child not a Cast; only rewire when safe.

-    for node in graph.node:
-        if is_pre_quant_scale_node(node):
-            pqs_child_nodes = [n for n in graph.node if node.output[0] in n.input]
-            assert len(pqs_child_nodes) == 1, f"Expected exactly one child node for {node.name}"
-            cast_node = pqs_child_nodes[0]
-            assert cast_node.op_type == "Cast", f"Expected Cast node for {node.name}"
-            node.output.clear()
-            node.output.extend(cast_node.output)
-            nodes_to_remove.append(cast_node.name)
+    for node in graph.node:
+        if is_pre_quant_scale_node(node):
+            pqs_child_nodes = [n for n in graph.node if node.output[0] in n.input]
+            if len(pqs_child_nodes) != 1 or pqs_child_nodes[0].op_type != "Cast":
+                continue
+            cast_node = pqs_child_nodes[0]
+            node.output[:] = cast_node.output
+            nodes_to_remove.append(cast_node.name)

---

1335-1343: Avoid magic number for BF16 detection

Use the enum for readability and safety.

-        for initializer in graph.initializer:
-            if initializer.data_type == 16:
+        for initializer in graph.initializer:
+            if initializer.data_type == onnx.TensorProto.BFLOAT16:
                 precision_dtype = "BFloat16"
                 break

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

modelopt/onnx/quantization/qdq_utils.py (1)

932-935: Reshape shape Constant detection is brittle (re-raising prior feedback)

String-matching "Constant" in input names is fragile. Use Reshape’s second input and check its producer.

-        # Remove constant node from reshape node
-        shape_constant_name = next(input for input in reshape_node.input if "Constant" in input)
-        nodes_to_remove.append(tensor_producer_map[shape_constant_name].name)
+        # Remove Constant that feeds Reshape's shape, if present
+        if len(reshape_node.input) >= 2:
+            shape_name = reshape_node.input[1]
+            shape_producer = tensor_producer_map.get(shape_name)
+            if shape_producer is not None and shape_producer.op_type == "Constant":
+                nodes_to_remove.append(shape_producer.name)

tests/unit/onnx/test_qdq_utils.py (1)

315-342: FP4 tests: packed shape and dtype OK; add an odd-first-dim guard test

• Expecting a packed shape along the first axis (2xN → 1xN) and dtype uint8 matches _cast_fp4’s contract.
• Consider adding a negative test where the first dim is odd to assert the error path.

I can add a parametric test that asserts the raised error for odd first-dimension inputs. Want me to push it?

Also applies to: 348-348

modelopt/onnx/quantization/qdq_utils.py (1)

1223-1229: Confirm FP4 initializer dims align with packing axis

Dims double the first axis, consistent with _cast_fp4’s first-axis packing. Verify all consumers assume this convention; earlier feedback suggested last-axis packing—ensure consistency across exporters/importers.

If you plan to switch packing to the last axis later, centralize the “packed axis” in one utility to avoid mismatches.

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

tests/_test_utils/import_helper.py (1)

80-92: Allow safe module-level skipping.

If this helper is ever used at import time, current skips lack allow_module_level=True. Harmless inside tests, but brittle at module level.

-    except importlib.metadata.PackageNotFoundError:
-        pytest.skip(f"{package_name} is not installed")
+    except importlib.metadata.PackageNotFoundError:
+        pytest.skip(f"{package_name} is not installed", allow_module_level=True)

-    if version.parse(installed_version) < version.parse(required_version):
-        pytest.skip(
-            f"{package_name} version {installed_version} is less than required {required_version}"
-        )
+    if version.parse(installed_version) < version.parse(required_version):
+        pytest.skip(
+            f"{package_name} version {installed_version} is less than required {required_version}",
+            allow_module_level=True,
+        )

tests/gpu/onnx/test_quantize_onnx_torch_int4_awq.py (2)

43-44: Gate ONNX version at import time to avoid ImportError before skip.

Current skip runs inside the test; onnx-dependent imports above may fail first in envs with <1.19 or missing onnx. Move the gate to module level and drop per-test calls.

-    skip_if_onnx_version_below_1_19()

Add near the imports (example):

# Top-level gate (place after standard imports)
from _test_utils.import_helper import skip_if_onnx_version_below_1_19
skip_if_onnx_version_below_1_19()

Note: This pairs with adding allow_module_level=True in the helper as suggested.

---

119-120: Same as above: move version gate to module level.

Prevents import-time failures and keeps skip reason centralized.

-    skip_if_onnx_version_below_1_19()

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

tests/_test_utils/import_helper.py (1)

16-20: Imports look fine; optional fallback for older Pythons.

If Python <3.8 is still in scope anywhere, add a fallback to importlib_metadata backport. Otherwise, keep as-is.

Apply if needed:

+try:
+    import importlib.metadata as importlib_metadata
+except ModuleNotFoundError:  # pragma: no cover
+    import importlib_metadata  # type: ignore

tests/gpu/onnx/test_quantize_onnx_torch_int4_awq.py (1)

42-44: Consider skipping at collection time with a decorator to reduce duplication.

Optional: switch to @pytest.mark.skipif(...) (using a small predicate in the helper) so tests are marked-skipped at collection, not runtime. Keeps both tests consistent and avoids executing setup code before the skip.

Example approach:

- def test_int4_awq(tmp_path):
-     skip_if_onnx_version_above_1_18()
+ @pytest.mark.skipif(_onnx_gt_1_18(), reason="Unsupported for ONNX > 1.18.0")
+ def test_int4_awq(tmp_path):
     ...
- def test_int4_awq_cuda(tmp_path):
-     skip_if_onnx_version_above_1_18()
+ @pytest.mark.skipif(_onnx_gt_1_18(), reason="Unsupported for ONNX > 1.18.0")
+ def test_int4_awq_cuda(tmp_path):

Helper (in import_helper.py):

def _onnx_gt_1_18() -> bool:
    try:
        v = importlib.metadata.version("onnx")
    except importlib.metadata.PackageNotFoundError:
        return False
    return version.parse(v) > version.parse("1.18.0")

Also applies to: 118-121

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro